#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <common/macro/build.h>
#include <common/macro/debug.h>
#include <common/avltree.h>
#include <common/basemacro.h>

/** 
* @brief Perform right rotation of the tree, with given relative root.
* 
* @param avltree Tree to be rotated.
* @param node Relative root.
*/
API_FILELOCAL void avltree_rotate_right(avltree_t *avltree, avltree_node_t *node) {
	avltree_node_t *tmp;

	assert(avltree != NULL);
	assert(node != NULL);

	node->left->parent = node->parent;
	if (node->parent != NULL) {
		if (node->parent->left == node)
			node->parent->left = node->left;
		else
			node->parent->right = node->left;
	}
	node->parent = node->left;

	tmp = node->left->right;
	node->left->right = node;
	node->left = tmp;
	if (node->left != NULL)
		node->left->parent = node;

	node->depth = MAX(AVLTREE_NODE_HEIGHT(node->left), AVLTREE_NODE_HEIGHT(node->right)) + 1;
	node->parent->depth = MAX(AVLTREE_NODE_HEIGHT(node->parent->left), AVLTREE_NODE_HEIGHT(node->parent->right)) + 1;

	if (node == avltree->root)
		avltree->root = node->parent;
}

/** 
* @brief Perform left rotation of the tree, with given relative root.
* 
* @param avltree Tree to be rotated.
* @param node Relative root.
*/
API_FILELOCAL void avltree_rotate_left(avltree_t *avltree, avltree_node_t *node) {
	avltree_node_t *tmp;

	assert(avltree != NULL);
	assert(node != NULL);

	node->right->parent = node->parent;
	if (node->parent != NULL) {
		if (node->parent->left == node)
			node->parent->left = node->right;
		else
			node->parent->right = node->right;
	}
	node->parent = node->right;

	tmp = node->right->left;
	node->right->left = node;
	node->right = tmp;
	if (node->right != NULL)
		node->right->parent = node;

	node->depth = MAX(AVLTREE_NODE_HEIGHT(node->left), AVLTREE_NODE_HEIGHT(node->right)) + 1;
	node->parent->depth = MAX(AVLTREE_NODE_HEIGHT(node->parent->left), AVLTREE_NODE_HEIGHT(node->parent->right)) + 1;

	if (node == avltree->root)
		avltree->root = node->parent;
}

/** 
* @brief Find predecessor of the given node in the its subtree.
* 
* @param avltree Tree in which we are looking for predecessor.
* @param node Relative root.
* 
* @return Predecessor or NULL.
*/
API_FILELOCAL avltree_node_t *avltree_subtree_predecessor(avltree_t *avltree, avltree_node_t *node) {
	avltree_node_t *replacement, *node_prev;

	assert(avltree != NULL);
	assert(node != NULL);

	replacement = node->left;
	node_prev = node;
	while (replacement != NULL) {
		node_prev = replacement;
		replacement = replacement->right;
	}
	if (node_prev != node)
		return node_prev;
	return NULL;
}

/** 
* @brief Find predecessor of the given node in the its supertree.
* 
* @param avltree Tree in which we are looking for predecessor.
* @param node Relative root.
* 
* @return Predecessor or NULL.
*/
API_FILELOCAL avltree_node_t *avltree_supertree_predecessor(avltree_t *avltree, avltree_node_t *node) {
	avltree_node_t *node_prev;

	assert(avltree != NULL);
	assert(node != NULL);

	node_prev = node;
	node = node->parent;
	while (node != NULL) {
		if (node->right == node_prev)
			return node;
		node_prev = node;
		node = node->parent;
	}

	return node;
}

/** 
* @brief Find successor of the given node in the its subtree.
* 
* @param avltree Tree in which we are looking for successor.
* @param node Relative root.
* 
* @return Successor or NULL.
*/
API_FILELOCAL avltree_node_t *avltree_subtree_successor(avltree_t *avltree, avltree_node_t *node) {
	avltree_node_t *replacement, *node_prev;

	assert(avltree != NULL);
	assert(node != NULL);

	replacement = node->right;
	node_prev = node;
	while (replacement != NULL) {
		node_prev = replacement;
		replacement = replacement->left;
	}
	if (node_prev != node)
		return node_prev;
	return NULL;
}

/** 
* @brief Find successor of the given node in the its supertree.
* 
* @param avltree Tree in which we are looking for successor.
* @param node Relative root.
* 
* @return Successor or NULL.
*/
API_FILELOCAL avltree_node_t *avltree_supertree_successor(avltree_t *avltree, avltree_node_t *node) {
	avltree_node_t *node_prev;

	assert(avltree != NULL);
	assert(node != NULL);

	node_prev = node;
	node = node->parent;
	while (node != NULL) {
		if (node->left == node_prev)
			return node;
		node_prev = node;
		node = node->parent;
	}

	return node;
}

/** 
* @brief Ensures correct balancing of the AVL tree, using L/R/RL/LR rotations.
* 
* @param avltree The tree to be rotated.
* @param node Source node for rebalancing, from which it is traversed up.
*/
API_FILELOCAL void avltree_rebalance(avltree_t *avltree, avltree_node_t *node) {
	int cmp;

	assert(avltree != NULL);
	assert(node != NULL);

	while (node != NULL) {
		/* Re-evaluate */
		node->depth = MAX(AVLTREE_NODE_HEIGHT(node->left), AVLTREE_NODE_HEIGHT(node->right)) + 1;

		cmp = AVLTREE_NODE_HEIGHT(node->left) - AVLTREE_NODE_HEIGHT(node->right);
		if (cmp >= 2) {
			cmp = AVLTREE_NODE_HEIGHT(node->left->left) - AVLTREE_NODE_HEIGHT(node->left->right);
			if (cmp == 1) {
				avltree_rotate_right(avltree, node);
			} else if (cmp == -1) {
				avltree_rotate_left(avltree, node->left);	
				avltree_rotate_right(avltree, node);
			}
		} else if (cmp <= -2) {
			cmp = AVLTREE_NODE_HEIGHT(node->right->left) - AVLTREE_NODE_HEIGHT(node->right->right);
			if (cmp == -1) {
				avltree_rotate_left(avltree, node);
			} else if (cmp == 1) {
				avltree_rotate_right(avltree, node->right);	
				avltree_rotate_left(avltree, node);	
			}
		}

		node = node->parent;
	}
}

/** 
* @brief Create new AVL tree.
* 
* @param cmp Comparator function, for key lookup in the tree
* 
* @return New avltree_t.
*/
API_LIBLOCAL avltree_t *avltree_new(generic_cmp_f cmp) {
	avltree_t *avltree;

	avltree = malloc(sizeof(*avltree));
	avltree->root = NULL;
	avltree->nodes = 0;

	avltree->cmp = cmp;

	return avltree;
}

/** 
* @brief Create new AVL tree node.
* 
* @param key Key of the node.
* @param value Value of the node.
* 
* @return New node.
*/
API_FILELOCAL avltree_node_t *avltree_node_new(void *key, void *value) {
	avltree_node_t *node;

	node = malloc(sizeof(*node));

	node->key = key;
	node->value = value;

	node->depth = 0;
	node->left = NULL;
	node->right = NULL;
	node->parent = NULL;

	return node;
}

/** 
* @brief Deallocate memory associated with the AVL tree.
* 
* @param avltree AVL tree structure.
*/
API_LIBLOCAL void avltree_destroy(avltree_t *avltree) {
	if (avltree == NULL) return;

	while (avltree->root != NULL) {
		avltree_delete(avltree, avltree->root);
	}
	free(avltree);
}

/** 
* @brief Deallocate memory associated with the AVL tree node.
* 
* @param node AVL tree node structure.
*/
API_FILELOCAL void avltree_node_destroy(avltree_node_t *node) {
	if (node == NULL) return;
	free(node);
}

/** 
* @brief Inserts a new node into the tree, and ensures proper balance.
* 
* @param avltree
* @param key
* @param value
* 
* @return The newly created AVL tree node structure.
*/
API_LIBLOCAL avltree_node_t *avltree_insert(avltree_t *avltree, void *key, void *value) {
	avltree_node_t *node, *node_prev, *result;
	int cmp;

	assert(avltree != NULL);

	node_prev = NULL;
	cmp = 0;

	/* We have to find where the node should go. */
	node = avltree->root;
	while (node != NULL) {
		node_prev = node;

		cmp = avltree->cmp(key, node->key);
		if (cmp > 0) {
			node = node->right;
		} else if (cmp < 0) {
			node = node->left;
		} else {
			/* Already in the tree. */
			return NULL;
		}
	}

	result = avltree_node_new(key, value);
	/* If avltree->root != NULL, then result_prev != NULL*/
	if (avltree->root != NULL) {
		result->parent = node_prev;
		if (cmp > 0) {
			result->parent->right = result;
		} else if (cmp < 0) {
			result->parent->left = result;
		}
	
	} else {
		avltree->root = result;
		result->depth = 0;
	}

	/* Node is in, now let's redepth all the way to the root. */
	if (result->parent != NULL)
		avltree_rebalance(avltree, result->parent);

	avltree->nodes++;

	return result;
}

/** 
* @brief Lookup a node according to the given key. Uses the comparator
* function passed in during AVL tree creation.
* 
* @param avltree
* @param key
* 
* @return 
*/
API_LIBLOCAL avltree_node_t *avltree_search(avltree_t *avltree, void *key) {
	avltree_node_t *node;
	int cmp;

	assert(avltree != NULL);

	node = avltree->root;
	while (node != NULL) {
		cmp = avltree->cmp(key, node->key);
		if (cmp > 0) {
			node = node->right;
		} else if (cmp < 0) {
			node = node->left;
		} else {
			return node;
		}
	}

	return NULL;
}


/** 
* @brief Returns the value associated with the given 
* 
* @param node
* 
* @return 
*/
API_LIBLOCAL void *avltree_node_value(avltree_node_t *node) {
	assert(node != NULL);
	return node->value;
}

/** 
* @brief Returns the key assocaited with the given node.
* 
* @param node
* 
* @return 
*/
API_LIBLOCAL void *avltree_node_key(avltree_node_t *node) {
	assert(node != NULL);
	return node->key;
}

/** 
* @brief Returns number of nodes in the AVL tree.
* 
* @param avltree
* 
* @return 
*/
API_LIBLOCAL int avltree_nodes(avltree_t *avltree) {
	assert(avltree != NULL);
	return avltree->nodes;
}

/** 
* @brief Find predecessor of the given node.
* 
* @param avltree Tree in which we are looking for predecessor.
* @param node
* 
* @return Predecessor or NULL.
*/
API_LIBLOCAL avltree_node_t *avltree_predecessor(avltree_t *avltree, avltree_node_t *node) {
	avltree_node_t *result;

	assert(avltree != NULL);
	assert(node != NULL);

	result = avltree_subtree_predecessor(avltree, node);
	if (result != NULL)
		return result;

	result = avltree_supertree_predecessor(avltree, node);
	return result;
}

/** 
* @brief Find successor of the given node.
* 
* @param avltree Tree in which we are looking for successor.
* @param node
* 
* @return Successor or NULL.
*/
API_LIBLOCAL avltree_node_t *avltree_successor(avltree_t *avltree, avltree_node_t *node) {
	avltree_node_t *result;

	assert(avltree != NULL);
	assert(node != NULL);

	result = avltree_subtree_successor(avltree, node);
	if (result != NULL)
		return result;

	result = avltree_supertree_successor(avltree, node);
	return result;
}

/** 
* @brief Deleted a node from the AVL tree, and ensures proper balance.
* 
* @param avltree
* @param node
* 
* @return 
*/
API_LIBLOCAL int avltree_delete(avltree_t *avltree, avltree_node_t *node) {
	avltree_node_t *replacement;

	assert(avltree != NULL);
	assert(node != NULL);

	if (node->left != NULL && node->right != NULL) {
		if (node->left->depth > node->right->depth) {
			replacement = avltree_subtree_predecessor(avltree, node);		
		} else {
			replacement = avltree_subtree_successor(avltree, node);
		}
	} else {
		replacement = avltree_subtree_predecessor(avltree, node);
		if (replacement == NULL)
			replacement = avltree_subtree_successor(avltree, node);
	}

	if (node->parent != NULL) {
		if (node->parent->left == node)
			node->parent->left = replacement;
		else
			node->parent->right = replacement;
	} else {
		avltree->root = replacement;
	}

	if (replacement != NULL) {
		if (replacement->parent != NULL) { 
			if (replacement == replacement->parent->left) {
				if (replacement->left != NULL) {
					replacement->parent->left = replacement->left;
					replacement->left->parent = replacement->parent;
				} else if (replacement->right != NULL) {
					replacement->parent->left = replacement->right;
					replacement->right->parent = replacement->parent;
				} else {
					replacement->parent->left = NULL;
				}
			} else {
				if (replacement->left != NULL) {
					replacement->parent->right = replacement->left;
					replacement->left->parent = replacement->parent;
				} else if (replacement->right != NULL) {
					replacement->parent->right = replacement->right;
					replacement->right->parent = replacement->parent;
				} else {
					replacement->parent->right = NULL;
				}
			}
		}

		replacement->parent = node->parent;
		replacement->depth = node->depth;

		if (node->left != replacement) {
			replacement->left = node->left;
		}
		if (node->right != replacement) {
			replacement->right = node->right;
		}

		if (replacement->left != NULL)
			replacement->left->parent = replacement;
		if (replacement->right != NULL)
			replacement->right->parent = replacement;

		avltree_rebalance(avltree, replacement);
	}

	avltree_node_destroy(node);
	avltree->nodes--;

	return 0;
}
